# -*- coding: utf-8 -*-
"""
   File Name：     ocr.py
   Description :  Extract car licenses in one image.
   Author :       KangJiaHui
   date：         2021/03/24
"""
import re
from paddleocr import PaddleOCR, draw_ocr
import cv2

model = PaddleOCR(use_angle_cls=True, lang="ch")


def car_plate_extract(image, coord):
    """
    Extract all car plates in input image.
    :param coord: list, [xmin, xmax, ymin, ymax], number means portion. e.x. [0, 0.5, 0.5, 1]
    :param image: numpy.ndarray
    :return: list,
    e.x. [["鲁FBR932", [[10,11], [25,11], [25,20], [10,20]]], ["鲁ABC123", [[20,21], [35,21], [35,30], [20,30]]]]
    """
    try:
        width = image.shape[1]
        height = image.shape[0]
        xmin = int(width * coord[0])
        xmax = int(width * coord[1])
        ymin = int(height * coord[2])
        ymax = int(height * coord[3])
        image_crop = image[ymin: ymax, xmin: xmax]
        ocr_result = model.ocr(image_crop, cls=True)
        print(ocr_result)
        plate = []
        for line in ocr_result:
            info = re.search('^[A-Z0-9]{4}[挂桂楼]$', line[-1][0])
            if info:
                bbox = line[0]
                info_2 = info.group(0)[:4] + "挂"
                plate.append([info_2, bbox])
        # 显示结果
        from PIL import Image
        image = Image.fromarray(cv2.cvtColor(image_crop, cv2.COLOR_BGR2RGB))
        boxes = [line[0] for line in ocr_result]
        txts = [line[1][0] for line in ocr_result]
        scores = [line[1][1] for line in ocr_result]
        im_show = draw_ocr(image, boxes, txts, scores, font_path='simfang.ttf')
        im_show = Image.fromarray(im_show)
        im_show.save('result001.jpg')
        return plate
    except Exception as e:
        msg = str(e)
        print(msg)
